package main

import (
	"fmt"
	"github.com/miekg/dns"
	"regexp"
	"strings"
	"time"
)

type ResolvError struct {
	qname       string
	nameservers []string
}

func (e ResolvError) Error() string {
	errmsg := fmt.Sprintf("%s resolv failed on %s", e.qname, strings.Join(e.nameservers, "; "))
	return errmsg
}

type Resolver struct {
	config        *dns.ClientConfig
	specResolvers []SpecResolver
	//map[string]SpecServerSettings
}

func (r *Resolver) Lookup(net string, req *dns.Msg) (message *dns.Msg, err error) {
	c := &dns.Client{
		Net:          net,
		ReadTimeout:  r.Timeout(),
		WriteTimeout: r.Timeout(),
	}

	qname := req.Question[0].Name
	Debug("qname:", qname)

	//自定义查询的示例代码
	i := 1
	for d, specDnsServer := range r.specResolvers {
		Debug(":Lookup specResolvers:", i, d, specDnsServer.specdomain, specDnsServer.dnsserver, specDnsServer.port, specDnsServer.search)
		var match bool = false
		switch specDnsServer.search {
		case "full":
			match = strings.Contains(qname, specDnsServer.specdomain)
			//Debug("full match:", match, qname, specDnsServer.specdomain)
		case "regex":
			match, _ = regexp.MatchString(specDnsServer.specdomain, qname)
			//Debug("full match:", match, qname, specDnsServer.specdomain)
		case "suffix":
			//Debug("here begin suffix match:", match, qname, specDnsServer.specdomain)
			fallthrough
		default:
			match = strings.HasSuffix(qname, specDnsServer.specdomain)
			//Debug("suffix match:", match, qname, specDnsServer.specdomain)
		}
		i = i + 1
		if match == false {
			//Debug("Don't match:", qname, specDnsServer.specdomain)
			continue
		}

		ns := specDnsServer.dnsserver + ":" + specDnsServer.port
		Debug("Begin Exchange:", qname, ns, specDnsServer.specdomain)
		r, rtt, err := c.Exchange(req, ns)
		if err != nil {
			Debug("%s socket error on %s", qname, ns)
			Debug("error:%s", err.Error())
			break
		}
		if r != nil && r.Rcode != dns.RcodeSuccess {
			Debug("%s failed to get an valid answer on %s", qname, ns)
			break
		}
		Debug("%s resolv on %s ttl: %d", UnFqdn(qname), ns, rtt)
		Debug("Reply:%s", r.String())

		return r, nil
	}
	//自定义查询的示例代码结尾

	for _, nameserver := range r.Nameservers() {
		Debug(":Lookup :", i, qname, nameserver)
		r, rtt, err := c.Exchange(req, nameserver)
		if err != nil {
			Debug("%s socket error on %s", qname, nameserver)
			Debug("error:%s", err.Error())
			continue
		}
		if r != nil && r.Rcode != dns.RcodeSuccess {
			Debug("%s failed to get an valid answer on %s", qname, nameserver)
			continue
		}
		Debug("%s resolv on %s ttl: %d", UnFqdn(qname), nameserver, rtt)
		Debug("Reply:%s", r.String())
		return r, nil
	}
	return nil, ResolvError{qname, r.Nameservers()}

}

func (r *Resolver) Nameservers() (ns []string) {
	for _, server := range r.config.Servers {
		nameserver := server + ":" + r.config.Port
		Debug("Resolver:", nameserver)
		ns = append(ns, nameserver)
	}
	return
}

func (r *Resolver) Timeout() time.Duration {
	return time.Duration(r.config.Timeout) * time.Second
}
